module wavefield_calculation
  use parameters
  use matrix_operations
  use grt_coefficients
  implicit none

contains

  ! 位移计算子程序
  subroutine funval(k, o, integ)
    use parameters, only : media, nly, cpn, csn, lo, ls, source
    real(kind=8), intent(in) :: k
    complex(kind=8), intent(in) :: o
    complex(kind=8), intent(out) :: integ(10)
    
    complex(kind=8) :: e(4,4), a0(2), a1(2), a2(2), b1, b2, fun(14)
    complex(kind=8) :: yu0, yd0, yu(2,2), yd(2,2)
    complex(kind=8) :: su0_1, su0_2, sd0_1, sd0_2
    complex(kind=8) :: su_0(2), su_1(2), su_2(2)
    complex(kind=8) :: sd_0(2), sd_1(2), sd_2(2)
    complex(kind=8) :: cigma0(2), cigma1(2), cigma2(2)
    complex(kind=8) :: cigma1_sh, cigma2_sh
    complex(kind=8) :: exu(2), exd(2)
    complex(kind=8) :: grdu(2,2), grud(2,2)
    complex(kind=8) :: grdu0, grud0
    complex(kind=8) :: gtu(2,2), gtu0
    complex(kind=8) :: csns, cpns, xc, vss, vps
    complex(kind=8) :: wru_0(2,2), wrd_0(2,2), wru_1(2,2), wrd_1(2,2)
    complex(kind=8) :: wru_2(2,2), wrd_2(2,2)
    real(kind=8) :: mus
    integer :: lay, i, j
    
    ! 确保复波数数组已分配
    if (.not. allocated(cpn)) then
      allocate(cpn(nly))
    end if
    if (.not. allocated(csn)) then
      allocate(csn(nly))
    end if
    
    ! 计算复波数
    do lay = 1, nly
      cpn(lay) = sqrt(k*k - (o/media%vp(lay))**2)
      csn(lay) = sqrt(k*k - (o/media%vs(lay))**2)
      if (real(cpn(lay)) < 0.0d0) cpn(lay) = -cpn(lay)
      if (real(csn(lay)) < 0.0d0) csn(lay) = -csn(lay)
    end do

    ! 计算反射/透射系数
    call grt_coefs(k, o)
    
    ! 震源计算
    exd(1) = exp(-cpn(ls)*(media%z(ls)-source%zs))
    exd(2) = exp(-csn(ls)*(media%z(ls)-source%zs))
    exu(1) = exp(-cpn(ls)*(source%zs-media%z(ls-1)))
    exu(2) = exp(-csn(ls)*(source%zs-media%z(ls-1)))
    
    csns = csn(ls)
    cpns = cpn(ls)
    vss = media%vs(ls)
    vps = media%vp(ls)
    mus = media%mu(ls)
    xc = 2.0d0 * vps * cpns * csns * mus * o / vss
    
    ! 震源向量计算
    cigma0(1) = vss * cpns * cpns * csns / xc
    cigma0(2) = -k * vps * cpns * csns / xc
    cigma1(1) = -2.0d0 * vss * csns * cpns * k / xc
    cigma1(2) = vps * (k*k + csns*csns) * cpns / xc
    cigma2(1) = vss * k*k * csns / xc
    cigma2(2) = -vps * cpns * csns * k / xc
    cigma2_sh = k / (2.0d0 * mus * csns)
    cigma1_sh = 1.0d0 / (2.0d0 * mus)
    
    ! 获取GR/GT系数
    do j = 1, 2
      do i = 1, 2
        grdu(i,j) = rt_coef%grdu(ls,i,j)
        grud(i,j) = rt_coef%grud(ls-1,i,j)
        gtu(i,j) = rt_coef%gtu(ls,i,j)
      end do
    end do
    grdu0 = rt_coef%grdu0(ls)
    grud0 = rt_coef%grud0(ls-1)
    gtu0 = rt_coef%gtu0(ls)
    
    do j = 1, 2
      do i = 1, 2
        wru_0(i,j) = exd(i) * grdu(i,j) * exd(j) + unit(i,j)
        wrd_0(i,j) = unit(i,j) + exu(i) * grud(i,j) * exu(j)
        wru_1(i,j) = exd(i) * grdu(i,j) * exd(j) - unit(i,j)
        wrd_1(i,j) = unit(i,j) - exu(i) * grud(i,j) * exu(j)
      end do
    end do
    
    su0_2 = (1.0d0 + exd(2) * grdu0 * exd(2)) * cigma2_sh
    sd0_2 = (exu(2) * grud0 * exu(2) + 1.0d0) * cigma2_sh
    su0_1 = (1.0d0 - exd(2) * grdu0 * exd(2)) * cigma1_sh
    sd0_1 = (exu(2) * grud0 * exu(2) - 1.0d0) * cigma1_sh
    
    do i = 1, 2
      su_0(i) = wru_0(i,1) * cigma0(1) + wru_0(i,2) * cigma0(2)
      sd_0(i) = wrd_0(i,1) * cigma0(1) + wrd_0(i,2) * cigma0(2)
      su_1(i) = wru_1(i,1) * cigma1(1) + wru_1(i,2) * cigma1(2)
      sd_1(i) = wrd_1(i,1) * cigma1(1) + wrd_1(i,2) * cigma1(2)
      su_2(i) = wru_0(i,1) * cigma2(1) + wru_0(i,2) * cigma2(2)
      sd_2(i) = wrd_0(i,1) * cigma2(1) + wrd_0(i,2) * cigma2(2)
    end do
    
    call mtxe(lo, k, o, e)
    ! Y矩阵计算
    call ydumtx(e, k, o, yu0, yd0, yu, yd)
    
    ! Uko计算
    if ((lo == ls .and. source%z0 <= source%zs) .or. lo < ls) then
      b1 = yu0 * su0_1
      b2 = yu0 * su0_2
      do i = 1, 2
        a0(i) = yu(i,1) * su_0(1) + yu(i,2) * su_0(2)
        a1(i) = yu(i,1) * su_1(1) + yu(i,2) * su_1(2)
        a2(i) = yu(i,1) * su_2(1) + yu(i,2) * su_2(2)
      end do
    else if ((lo == ls .and. source%z0 > source%zs) .or. lo > ls) then
      b1 = yd0 * sd0_1
      b2 = yd0 * sd0_2
      do i = 1, 2
        a0(i) = yd(i,1) * sd_0(1) + yd(i,2) * sd_0(2)
        a1(i) = yd(i,1) * sd_1(1) + yd(i,2) * sd_1(2)
        a2(i) = yd(i,1) * sd_2(1) + yd(i,2) * sd_2(2)
      end do
    end if
    
    ! 计算最终结果
    fun(1) = -b1
    fun(2) = -a2(1)
    fun(3) = a0(1)
    fun(4) = b2
    fun(5) = (b1 - a1(1))
    fun(6) = (a2(1) - b2)
    fun(7) = a1(1)
    fun(8) = a2(1)
    fun(9) = fun(5)
    fun(10) = -fun(6)
    fun(11) = a2(2)
    fun(12) = -a0(2)
    fun(13) = a1(2)
    fun(14) = -a2(2)
    
    integ(1) = fun(7) + fun(9)
    integ(2) = fun(8) + fun(10)
    integ(3) = fun(4) + fun(6)
    integ(4) = fun(1) + fun(5)
    integ(5) = fun(2)
    integ(6) = fun(3)
    integ(7) = fun(14)
    integ(8) = fun(13)
    integ(9) = fun(11)
    integ(10) = fun(12)
  end subroutine funval

  ! Y矩阵计算子程序
  subroutine ydumtx(e, k, o, yu0, yd0, yu, yd)
    use parameters, only : media, nly, cpn, csn, lo, ls, source, rt_coef
    use matrix_operations, only : mtxinv_2, mtxmtp
    complex(kind=8), intent(in) :: e(4,4), o
    real(kind=8), intent(in) :: k
    complex(kind=8), intent(out) :: yu0, yd0, yu(2,2), yd(2,2)
    
    complex(kind=8) :: yu_psv(4,4), yd_psv(4,4), y_sh, temp1(4,4), temp2(4,4)
    complex(kind=8) :: exdd(2), exuu(2), exds(2),exus(2)
    complex(kind=8) :: u0, d0, u(2,2), d(2,2)
    complex(kind=8) :: exd(2), exu(2)
    complex(kind=8) :: grdu(2,2), grud(2,2)
    complex(kind=8) :: grdu0, grud0
    integer :: i, j, lay
    
    exd(1) = exp(-cpn(ls)*(media%z(ls)-source%zs))
    exd(2) = exp(-csn(ls)*(media%z(ls)-source%zs))
    exu(1) = exp(-cpn(ls)*(source%zs-media%z(ls-1)))
    exu(2) = exp(-csn(ls)*(source%zs-media%z(ls-1)))
    
    exdd(1) = exp(-cpn(lo)*(source%z0-media%z(lo-1)))
    exdd(2) = exp(-csn(lo)*(source%z0-media%z(lo-1)))
    
    if (lo == nly) then
      exuu(1) = 0.0d0
      exuu(2) = 0.0d0
    else
      exuu(1) = exp(-cpn(lo)*(media%z(lo)-source%z0))
      exuu(2) = exp(-csn(lo)*(media%z(lo)-source%z0))
    end if
    
    ! 获取GR/GT系数
    do j = 1, 2
      do i = 1, 2
        grdu(i,j) = rt_coef%grdu(ls,i,j)
        grud(i,j) = rt_coef%grud(ls-1,i,j)
      end do
    end do
    grdu0 = rt_coef%grdu0(ls)
    grud0 = rt_coef%grud0(ls-1)
    
    do j = 1, 2
      do i = 1, 2
        temp1(i,j) = exd(i) * grdu(i,j) * exd(j)
        temp2(i,j) = exu(i) * grud(i,j) * exu(j)
      end do
    end do

    call mtxmtp(2, temp1, temp2, yu_psv)
    call mtxmtp(2, temp2, temp1, yd_psv)

    do j = 1, 2
      do i = 1, 2
        yu_psv(i,j) = unit(i,j) - yu_psv(i,j)
        yd_psv(i,j) = unit(i,j) - yd_psv(i,j)
      end do
    end do
    call mtxinv_2(yu_psv)
    call mtxinv_2(yd_psv)
    
    y_sh = 1.0d0 / (1.0d0 - exd(2)*grdu0*exd(2)*exu(2)*grud0*exu(2))
    
    ! 根据接收器和源的相对位置计算
    if (lo == ls .and. source%z0 <= source%zs) then
      exds(1) = exp(-cpn(ls)*(source%z0-media%z(ls-1)))
      exds(2) = exp(-csn(ls)*(source%z0-media%z(ls-1)))
      exus(1) = exp(-cpn(ls)*(source%zs-source%z0))
      exus(2) = exp(-csn(ls)*(source%zs-source%z0))
      
      u0 = exds(2) * grud0 * exu(2) + exus(2)
      do j = 1, 2
        do i = 1, 2
          u(i,j) = e(i,1)*exds(1)*rt_coef%grud(ls-1,1,j)*exu(j) + &
                   e(i,2)*exds(2)*rt_coef%grud(ls-1,2,j)*exu(j) + &
                   e(i,j+2)*exus(j)
        end do
      end do
      
      yu0 = u0 * y_sh
      do j = 1, 2
        do i = 1, 2
          yu(i,j) = u(i,1)*yu_psv(1,j) + u(i,2)*yu_psv(2,j)
        end do
      end do
      
    else if (lo == ls .and. source%z0 > source%zs) then
      exds(1) = exp(-cpn(ls)*(source%z0-source%zs))
      exds(2) = exp(-csn(ls)*(source%z0-source%zs))
      exus(1) = exp(-cpn(ls)*(media%z(ls)-source%z0))
      exus(2) = exp(-csn(ls)*(media%z(ls)-source%z0))
      
      d0 = exds(2) + exus(2) * grdu0 * exd(2)
      do j = 1, 2
        do i = 1, 2
          d(i,j) = e(i,j)*exds(j) + &
                   e(i,3)*exus(1)*rt_coef%grdu(ls,1,j)*exd(j) + &
                   e(i,4)*exus(2)*rt_coef%grdu(ls,2,j)*exd(j)
        end do
      end do
      
      yd0 = d0 * y_sh
      do j = 1, 2
        do i = 1, 2
          yd(i,j) = d(i,1)*yd_psv(1,j) + d(i,2)*yd_psv(2,j)
        end do
      end do
      
    else if (lo < ls) then
      u0 = exdd(2) * rt_coef%grud0(lo-1) * rt_coef%ex(lo,2) + exuu(2)
      do j = 1, 2
        do i = 1, 2
          u(i,j) = e(i,1)*exdd(1)*rt_coef%grud(lo-1,1,j)*rt_coef%ex(lo,j) + &
                   e(i,2)*exdd(2)*rt_coef%grud(lo-1,2,j)*rt_coef%ex(lo,j) + &
                   e(i,j+2)*exuu(j)
        end do
      end do
      
      do lay = ls-1, lo, -1
        if (lay == ls-1) then
          yu0 = rt_coef%gtu0(lay) * exu(2) * y_sh
          do j = 1, 2
            do i = 1, 2
              yu(i,j) = rt_coef%gtu(lay,i,1)*exu(1)*yu_psv(1,j) + &
                        rt_coef%gtu(lay,i,2)*exu(2)*yu_psv(2,j)
            end do
          end do
        else
          yu0 = rt_coef%gtu0(lay) * rt_coef%ex(lay+1,2) * yu0
          do j = 1, 2
            do i = 1, 2
              temp1(i,j) = rt_coef%gtu(lay,i,1)*rt_coef%ex(lay+1,1)*yu(1,j) + &
                           rt_coef%gtu(lay,i,2)*rt_coef%ex(lay+1,2)*yu(2,j)
            end do
          end do
          do j = 1, 2
            do i = 1, 2
              yu(i,j) = temp1(i,j)
            end do
          end do
        end if
      end do
      
      yu0 = u0 * yu0
      do j = 1, 2
        do i = 1, 2
          temp1(i,j) = u(i,1)*yu(1,j) + u(i,2)*yu(2,j)
        end do
      end do
      
      do j = 1, 2
        do i = 1, 2
          yu(i,j) = temp1(i,j)
        end do
      end do
      
    else if (lo > ls) then
      d0 = exdd(2) + exuu(2) * rt_coef%grdu0(lo) * rt_coef%ex(lo-1,2)
      do j = 1, 2
        do i = 1, 2
          d(i,j) = e(i,j)*exdd(j) + &
                   e(i,3)*exuu(1)*rt_coef%grdu(lo,1,j)*rt_coef%ex(lo-1,j) + &
                   e(i,4)*exuu(2)*rt_coef%grdu(lo,2,j)*rt_coef%ex(lo-1,j)
        end do
      end do
      
      do lay = ls+1, lo
        if (lay == ls+1) then
          yd0 = rt_coef%gtd0(ls) * exd(2) * y_sh
          do j = 1, 2
            do i = 1, 2
              yd(i,j) = rt_coef%gtd(ls,i,1)*exd(1)*yd_psv(1,j) + &
                        rt_coef%gtd(ls,i,2)*exd(2)*yd_psv(2,j)
            end do
          end do
        else
          yd0 = rt_coef%gtd0(lay-1) * rt_coef%ex(lay-2,2) * yd0
          do j = 1, 2
            do i = 1, 2
              temp1(i,j) = rt_coef%gtd(lay-1,i,1)*rt_coef%ex(lay-2,1)*yd(1,j) + &
                           rt_coef%gtd(lay-1,i,2)*rt_coef%ex(lay-2,2)*yd(2,j)
            end do
          end do
          do j = 1, 2
            do i = 1, 2
              yd(i,j) = temp1(i,j)
            end do
          end do
        end if
      end do
      
      yd0 = d0 * yd0
      do j = 1, 2
        do i = 1, 2
          temp1(i,j) = d(i,1)*yd(1,j) + d(i,2)*yd(2,j)
        end do
      end do
      do j = 1, 2
        do i = 1, 2
          yd(i,j) = temp1(i,j)
        end do
      end do
    end if
  end subroutine ydumtx

end module wavefield_calculation